#! /usr/bin/env python
# coding=utf-8

import argparse
import sys
from argparse import RawTextHelpFormatter
import os
from Bio import SeqIO

parser = argparse.ArgumentParser(
    description='''
    读取mRNA 改名字 并且获得 可以翻译的mRNA 和 pep 序列 作为文件输出
    give_mRNA_change_name_and_make_pep.py -p Amel -r mRNA.fa -o match_lista -a changed_mRNA.fa -b changed_pep.fa

    ''',formatter_class=RawTextHelpFormatter)

parser.add_argument('-p',
                help='输出物种命名的前缀')

parser.add_argument('-r',
                help='原始 mRNA.fa 文件')

parser.add_argument('-o',
                help='输出配对列表')

parser.add_argument('-a',
                help='输出的mRNA文件')

parser.add_argument('-b',
                help='输出的pep文件')


args = parser.parse_args()


if not args.p or not args.r or not args.o or not args.a or not args.b:
    parser.print_help()
    sys.exit()




# 会生成一个对应表格

infile_mRNA = args.r

match_lista = open(args.o,'w')

outfile_mRNA = args.a
outfile_pep = args.b

outfile_mRNA = open(outfile_mRNA ,'w')

outfile_pep = open(outfile_pep,'w')

prefix = args.p



n = 1

for i in SeqIO.parse(infile_mRNA,'fasta'):
	raw_name = str(i.name)
	new_name = prefix + '_'+str(n)
	pep_seq = str(i.seq.translate())
	if pep_seq.strip('.').find('.')!=-1 or pep_seq.strip('*').find('*')!=-1:
		print('can not translate for %s' % raw_name)
		continue
	else:
		match_lista.write(raw_name+'\t'+new_name+'\n')
		outfile_pep.write('>'+new_name+'\n'+pep_seq+'\n')
		outfile_mRNA.write('>'+new_name+'\n'+str(i.seq)+'\n')
		n+=1


match_lista.close()
outfile_pep.close()
outfile_mRNA.close()